package org.xyl.repository.impl;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.xyl.service.embedding.SimpleEmbeddingService;
import org.springframework.stereotype.Service;
import org.xyl.config.VectorDbConfig;
import org.xyl.model.KnowledgeDocument;
import org.xyl.repository.VectorDbService;

import java.io.*;
import java.util.*;
import java.util.stream.Collectors;

/**
 * 内存向量数据库服务实现
 * 使用Spring AI的EmbeddingClient进行向量生成
 */
@Service
public class InMemoryVectorDbServiceImpl implements VectorDbService {
    
    private static final Logger logger = LoggerFactory.getLogger(InMemoryVectorDbServiceImpl.class);
    
    @Autowired
    private VectorDbConfig vectorDbConfig;
    
    @Autowired
    private SimpleEmbeddingService embeddingService;
    
    // 简化的向量存储结构
    private static class VectorEntry {
        private String id;
        private String content;
        private Map<String, Object> metadata;
        private List<Double> embedding;
        
        public VectorEntry(String id, String content, Map<String, Object> metadata, List<Double> embedding) {
            this.id = id;
            this.content = content;
            this.metadata = metadata;
            this.embedding = embedding;
        }
    }
    
    private final List<VectorEntry> vectorStore = new ArrayList<>();
    
    @Override
    public void initialize() {
        try {
            // 尝试从文件加载持久化数据
            loadFromPersistence();
            logger.info("向量数据库初始化成功");
        } catch (Exception e) {
            logger.error("向量数据库初始化失败", e);
            throw new RuntimeException("向量数据库初始化失败", e);
        }
    }
    
    @Override
    public void storeDocuments(List<KnowledgeDocument> documents) {
        if (documents == null || documents.isEmpty()) {
            return;
        }
        
        try {
            for (KnowledgeDocument document : documents) {
                // 生成嵌入向量
                List<Double> embedding = embeddingService.embed(document.getContent());
                
                // 存储向量条目
                VectorEntry entry = new VectorEntry(
                    document.getId(),
                    document.getContent(),
                    document.getMetadata() != null ? document.getMetadata() : new HashMap<>(),
                    embedding
                );
                
                vectorStore.add(entry);
                logger.debug("文档存储成功: {}", document.getId());
            }
            
            // 保存到持久化存储
            saveToPersistence();
            
            logger.info("成功存储 {} 个文档", documents.size());
        } catch (Exception e) {
            logger.error("文档存储失败", e);
            throw new RuntimeException("文档存储失败", e);
        }
    }
    
    @Override
    public List<KnowledgeDocument> search(String query, int topK, Map<String, Object> filters) {
        try {
            // 生成查询嵌入
            List<Double> queryEmbedding = embeddingService.embed(query);
            
            // 计算相似度并排序
            List<Map.Entry<VectorEntry, Double>> similarities = vectorStore.stream()
                .filter(entry -> matchesFilters(entry.metadata, filters))
                .map(entry -> {
                    double similarity = calculateCosineSimilarity(queryEmbedding, entry.embedding);
                    return new AbstractMap.SimpleEntry<>(entry, similarity);
                })
                .sorted(Map.Entry.<VectorEntry, Double>comparingByValue().reversed())
                .collect(Collectors.toList());
            
            // 转换为KnowledgeDocument
            List<KnowledgeDocument> results = new ArrayList<>();
            int limit = Math.min(topK, similarities.size());
            
            for (int i = 0; i < limit; i++) {
                Map.Entry<VectorEntry, Double> entry = similarities.get(i);
                VectorEntry vectorEntry = entry.getKey();
                Double similarity = entry.getValue();
                
                KnowledgeDocument document = new KnowledgeDocument();
                document.setId(vectorEntry.id);
                document.setContent(vectorEntry.content);
                document.setMetadata(new HashMap<>(vectorEntry.metadata));
                document.getMetadata().put("confidence", similarity);
                
                results.add(document);
            }
            
            return results;
        } catch (Exception e) {
            logger.error("文档搜索失败: {}", query, e);
            throw new RuntimeException("文档搜索失败", e);
        }
    }
    
    @Override
    public void deleteDocument(String documentId) {
        try {
            boolean removed = vectorStore.removeIf(entry -> entry.id.equals(documentId));
            if (removed) {
                logger.info("文档删除成功: {}", documentId);
                // 保存到持久化存储
                saveToPersistence();
            } else {
                logger.warn("文档未找到: {}", documentId);
            }
        } catch (Exception e) {
            logger.error("删除文档失败", e);
            throw new RuntimeException("删除文档失败", e);
        }
    }
    
    @Override
    public void clear() {
        try {
            vectorStore.clear();
            
            // 删除持久化文件
            File persistenceFile = new File(getPersistenceFilePath());
            if (persistenceFile.exists()) {
                persistenceFile.delete();
            }
            
            logger.info("向量数据库已清空");
        } catch (Exception e) {
            logger.error("清空向量数据库失败", e);
            throw new RuntimeException("清空向量数据库失败", e);
        }
    }
    
    // 计算余弦相似度
    private double calculateCosineSimilarity(List<Double> vec1, List<Double> vec2) {
        double dotProduct = 0.0;
        double norm1 = 0.0;
        double norm2 = 0.0;
        
        int minLength = Math.min(vec1.size(), vec2.size());
        
        for (int i = 0; i < minLength; i++) {
            dotProduct += vec1.get(i) * vec2.get(i);
            norm1 += vec1.get(i) * vec1.get(i);
            norm2 += vec2.get(i) * vec2.get(i);
        }
        
        if (norm1 == 0 || norm2 == 0) {
            return 0.0;
        }
        
        return dotProduct / (Math.sqrt(norm1) * Math.sqrt(norm2));
    }
    
    // 保存到持久化存储
    private void saveToPersistence() {
        try {
            // 确保目录存在
            File persistenceDir = new File(vectorDbConfig.getPersistencePath());
            if (!persistenceDir.exists()) {
                persistenceDir.mkdirs();
            }
            
            // 保存向量数据
            String filePath = getPersistenceFilePath();
            try (ObjectOutputStream oos = new ObjectOutputStream(new FileOutputStream(filePath))) {
                // 我们需要一个可序列化的结构
                List<SerializableVectorEntry> serializableEntries = vectorStore.stream()
                    .map(entry -> new SerializableVectorEntry(
                        entry.id,
                        entry.content,
                        new HashMap<>(entry.metadata),
                        new ArrayList<>(entry.embedding)
                    ))
                    .collect(Collectors.toList());
                
                oos.writeObject(serializableEntries);
            }
            
            logger.debug("数据已持久化到: {}", filePath);
        } catch (Exception e) {
            logger.warn("数据持久化失败", e);
        }
    }
    
    // 从持久化存储加载
    @SuppressWarnings("unchecked")
    private void loadFromPersistence() {
        try {
            String filePath = getPersistenceFilePath();
            File persistenceFile = new File(filePath);
            
            if (persistenceFile.exists()) {
                try (ObjectInputStream ois = new ObjectInputStream(new FileInputStream(filePath))) {
                    List<SerializableVectorEntry> serializableEntries = (List<SerializableVectorEntry>) ois.readObject();
                    
                    // 转换回VectorEntry
                    for (SerializableVectorEntry entry : serializableEntries) {
                        vectorStore.add(new VectorEntry(
                            entry.id,
                            entry.content,
                            entry.metadata,
                            entry.embedding
                        ));
                    }
                    
                    logger.info("从持久化存储加载了 {} 个文档", serializableEntries.size());
                }
            }
        } catch (Exception e) {
            logger.warn("从持久化存储加载数据失败", e);
        }
    }
    
    // 获取持久化文件路径
    private String getPersistenceFilePath() {
        return vectorDbConfig.getPersistencePath() + File.separator + "vector_db.ser";
    }
    
    // 检查是否匹配过滤条件
    private boolean matchesFilters(Map<String, Object> metadata, Map<String, Object> filters) {
        if (filters == null || filters.isEmpty()) {
            return true;
        }
        
        for (Map.Entry<String, Object> filter : filters.entrySet()) {
            String key = filter.getKey();
            Object value = filter.getValue();
            
            if (!metadata.containsKey(key)) {
                return false;
            }
            
            Object metadataValue = metadata.get(key);
            if (metadataValue == null || !metadataValue.toString().equals(String.valueOf(value))) {
                return false;
            }
        }
        
        return true;
    }
    
    // 可序列化的向量条目类
    private static class SerializableVectorEntry implements Serializable {
        private static final long serialVersionUID = 1L;
        
        private String id;
        private String content;
        private Map<String, Object> metadata;
        private List<Double> embedding;
        
        public SerializableVectorEntry(String id, String content, Map<String, Object> metadata, List<Double> embedding) {
            this.id = id;
            this.content = content;
            this.metadata = metadata;
            this.embedding = embedding;
        }
    }
}